__author__ = 'liyane'

import json
import os

# 初始化环境和OpenAI
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())
openai_api_base = os.environ['OPENAI_API_BASE']
openai_api_key = os.environ['OPENAI_API_KEY']

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base
)

#调用GPT大模型
def get_completion(messages, tools, model="Qwen/Qwen2.5-14B-Instruct"):
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        # tool_choice支持设置 "auto"（由模型决定是否调用tool） 或者 "none" （不调用tool）作为value。 有tools定义时默认由模型决定。
        # 也可以强制要求必须调用指定的函数，如下所示
        # tool_choice= {"type": "function", "function": {"name": "multiply"}} ,
        tools=tools
    )
    return response.choices[0].message

#定义function/tool 1: multiply
def multiply(first_int: int, second_int: int) -> int:
    """两个整数相乘"""
    return first_int * second_int

#定义function/tool 2: add
def add(first_add: int, second_add: int) -> int:
    """两个整数相加"""
    return first_add + second_add

#将function calling的schema格式告诉大模型
tools=[{
    "type": "function",
    "function": {
        "name": "multiply",
        "description": "两个整数相乘",
        "parameters": {
            "type": "object",
            "properties": {
                "first_int": {
                    "type": "integer",
                    "description": "第一个乘数",
                },
                "second_int": {
                    "type": "integer",
                    "description": "第二个乘数",
                }
            },
            "required": ["first_int", "second_int"],
        }
    }
    },
    {
    "type": "function",
    "function": {
        "name": "add",
        "description": "两个整数相加",
        "parameters": {
            "type": "object",
            "properties": {
                "first_add": {
                    "type": "integer",
                    "description": "第一个加数",
                },
                "second_add": {
                    "type": "integer",
                    "description": "第二个加数",
                }
            },
            "required": ["first_add", "second_add"],
        }
    }
}]

# 调用LLM接口，将LLM回复加入对话上下文
prompt = "一共有3个人，每个人有15个苹果，10个鸭梨，一共有多少苹果？有多少水果？,"
messages = [
    {"role": "user", "content": prompt}
]
response = get_completion(messages, tools)
messages.append(response)

# 如果LLM需要function calling，调用相应的函数，并将函数结果数据加入对话上下文，继续调用LLM。
while (response.tool_calls is not None):
    for tool_call in response.tool_calls:
        selected_tool = {"add": add, "multiply": multiply}[tool_call.function.name]
        args = json.loads(tool_call.function.arguments)
        tool_output = selected_tool(**args)

        messages.append({
            "tool_call_id": tool_call.id,  # 用于标识函数调用的 ID
            "role": "tool",
            "name": tool_call.function.name,
            "content": str(tool_output)  # 数值result 必须转成字符串
        })

    response = get_completion(messages, tools)
    messages.append(response)

print("=====最终回复=====")
print(response.content)

